from __future__ import with_statement

import sys, optparse, os.path

import numpy, scipy.stats

import common, stats
from retraction_unfolding_filter import CurveStretch, RetractionPeak

def statistics_string(s, x):
    if not len(x):
        return "%s: None" % (s, )
    x_mean, x_median, x_std = numpy.mean(x), numpy.median(x), numpy.std(x)
    return "%s: Count: %d  Mean: %s  Median: %s  Std: %s  Min: %s  Max: %s" % (
        s, len(x), numpy.mean(x), numpy.median(x), numpy.std(x), numpy.min(x), numpy.max(x))
    # gumbel_mode, gumbel_scale = stats.gumbel_l_loc_scale(x)
    # LeftGumbelMode: %s  LeftGumbelScale: %s" % (gumbel_mode, gumbel_scale)

def main(params):
    parser = optparse.OptionParser()
    parser.add_option("--processed_playlist", help="input .repr or .hkp file")
    parser.add_option(
        "--curve_stretch_playlist",
        help="optional input curve stretch playlist", default="")
    parser.add_option("--min_delta_cl_nm")
    parser.add_option("--max_delta_cl_nm")
    parser.add_option("--min_rupture_force_pn")
    parser.add_option("--max_rupture_force_pn")
    parser.add_option("--delta_force_filename", help="output filename")
    parser.add_option("--scatterplot_filename", help="eps filename")

    options, args = parser.parse_args(params)
    if args:
        parser.print_help()
        return -1
    print(options)
    min_delta_cl_nm = float(options.min_delta_cl_nm)
    max_delta_cl_nm = float(options.max_delta_cl_nm)
    min_rupture_force_pn = float(options.min_rupture_force_pn)
    max_rupture_force_pn = float(options.max_rupture_force_pn)
    scatterplot_filename = options.scatterplot_filename

    common.check_True(
        os.path.isfile(options.processed_playlist),
        "%s does not exist or it is not a file" % options.processed_playlist)
    # Since the curve stretch does not have the last peak, we also need .hkp.
    delta = []
    if os.path.isfile(options.curve_stretch_playlist):
        path2len = {}
        with open(options.curve_stretch_playlist, "rt") as f:
            for line in f:
                curve_stretch = eval(line)
                path2len[curve_stretch.path] = len(curve_stretch)
                # TODO: We should check compatibility between .hkp and curve_stretch.
        path2len_function = lambda path, n: n - path2len.get(path, 0)
    else:
        path2len_function = lambda path, n: 1
    with open(options.processed_playlist, "rt") as f:
        for line in f:
            path, data = stats.extract_path_data(line)
            s = path2len_function(path, len(data))
            for i in range(s, len(data)):
                d = data[i][0] - data[i - 1][0]
                delta.append((d, data[i - 1][1]))
    delta = [
        (d, f) for (d, f) in delta
        if min_delta_cl_nm <= d <= max_delta_cl_nm and min_rupture_force_pn <= f <= max_rupture_force_pn]
    
    with open(options.delta_force_filename, "wt") as f:
        for d in delta:
            f.write("%r, %r\n" % d)

    if scatterplot_filename:
        if scatterplot_filename.lower().endswith(".eps") or scatterplot_filename.lower().endswith(".ps"):
            x, y = zip(*delta) if delta else ([], [])
            print(statistics_string("Contour Length (nm)", x))
            print(statistics_string("Force (pN)", y))

            import matplotlib
            # Avoid the need for X11.
            matplotlib.use("PS")
            import pylab
            import matplotlib.backends.backend_ps
            import matplotlib.ticker
            import numpy
            pylab.rc("font", family="serif")
            pylab.rc("font", size=10)
            pylab.rc("axes", linewidth=0.5)
            pylab.rc('legend', fontsize='small')
            nullfmt = matplotlib.ticker.NullFormatter()

            # definitions for the axes
            left, width = 0.1, 0.65
            bottom, height = 0.1, 0.65
            bottom_h = left_h = left + width + 0.02
            rect_scatter = [left, bottom, width, height]
            rect_histx = [left, bottom_h, width, 0.2]
            rect_histy = [left_h, bottom, 0.2, height]

            figure = matplotlib.pyplot.figure(1, figsize=(8, 8))

            axScatter = matplotlib.pyplot.axes(rect_scatter)
            axHistx = matplotlib.pyplot.axes(rect_histx)
            axHisty = matplotlib.pyplot.axes(rect_histy)

            # no labels
            axHistx.xaxis.set_major_formatter(nullfmt)
            axHisty.yaxis.set_major_formatter(nullfmt)

            if len(x):
                # the scatter plot:
                axScatter.scatter(x, y, s=1, facecolors="none")

                # now determine nice limits by hand:
                xbinwidth, ybinwidth = 1, 10
                xymax = max( [max(x), max(y)] )

                axScatter.set_xlim(0, max(x) + xbinwidth)
                axScatter.set_ylim(0, max(y) + ybinwidth)

                _, xbins, _ = axHistx.hist(x, bins=numpy.arange(0, max(x) + xbinwidth, xbinwidth))
                xbincenters = 0.5 * (xbins[1:] + xbins[:-1])
                x_mean, x_std = numpy.mean(x), numpy.std(x)
                xpdf = scipy.stats.norm.pdf(xbincenters, x_mean, x_std)
                axHistx.plot(xbincenters, xpdf * len(x) * xbinwidth, 'r--', linewidth=1)
                axHistx.text(0, 0, "$\mu$=%.1f nm\nmedian=%.1f nm\n$\sigma$=%.1f nm" % (x_mean, numpy.median(x), x_std))

                _, ybins, _ = axHisty.hist(
                    y, bins=numpy.arange(0, max(y) + ybinwidth, ybinwidth), orientation='horizontal')
                ybincenters = 0.5 * (ybins[1:] + ybins[:-1])
                y_mean, y_std = numpy.mean(y), numpy.std(y)
                ypdf = scipy.stats.norm.pdf(ybincenters, y_mean, y_std)
                axHisty.plot(ypdf * len(y) * ybinwidth, ybincenters, 'r--', linewidth=1)
                axHisty.text(0, 0, "$\mu$=%.1f pN\nmedian=%.1f pN\n$\sigma$=%.1f pN" % (y_mean, numpy.median(y), y_std))

                axHistx.set_xlim(axScatter.get_xlim())
                axHisty.set_ylim(axScatter.get_ylim())

            axScatter.set_xlabel("delta_cl (after) [nm]")
            axScatter.set_ylabel("unfolding force [pN]")

            canvas = matplotlib.backends.backend_ps.FigureCanvasPS(figure)
            canvas.print_figure(scatterplot_filename)
        else:
            print("Unknown filename: %s" % scatterplot_filename)

if __name__ == "__main__":
    main(sys.argv[1:])
